from typing import Annotated
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_ollama import ChatOllama
import json
from langchain_core.messages import ToolMessage
from langchain_core.tools import InjectedToolCallId, tool
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.types import Command, interrupt
from langgraph.checkpoint.memory import MemorySaver

memory = MemorySaver()

class State(TypedDict):
    # Messages have the type "list". The `add_messages` function
    # in the annotation defines how this state key should be updated
    # (in this case, it appends messages to the list, rather than overwriting them)
    messages: Annotated[list, add_messages]
    name: str
    birthday: str

llm = ChatOllama(
    model="qwen2.5:7b-instruct-q5_K_S",
    temperature=0.2,
)

@tool
def get_birthday(person: str) -> str:
    """Get the birthday for a given person's name."""
    # print(f"Getting birthday for {person}")
    return f" {person}'s birthyday is 2106-03-03."

@tool
# Note that because we are generating a ToolMessage for a state update, we
# generally require the ID of the corresponding tool call. We can use
# LangChain's InjectedToolCallId to signal that this argument should not
# be revealed to the model in the tool's schema.
def human_assistance(name: str, birthday: str, tool_call_id: Annotated[str, InjectedToolCallId]) -> str:
    """Request assistance from a human."""
    # print(f"Requesting human assistance for: {query}")
    print(f"Is this correct? {name} and {birthday}. 输入'y'确认，或输入修正信息")
    human_response = interrupt(
        {
            "question": "Is this correct?",
            "name": name,
            "birthday": birthday,
        },
    )
    # If the information is correct, update the state as-is.
    print(f"<<<Human response>>>: {human_response}")
    if human_response.get("correct", "").lower().startswith("y"):
        verified_name = name
        verified_birthday = birthday
        response = "Correct"
    # Otherwise, receive information from the human reviewer.
    else:
        verified_name = human_response.get("name", name)
        verified_birthday = human_response.get("birthday", birthday)
        response = f"Made a correction: {human_response}"

    # This time we explicitly update the state with a ToolMessage inside
    # the tool.
    state_update = {
        "name": verified_name,
        "birthday": verified_birthday,
        "messages": [ToolMessage(response, tool_call_id=tool_call_id)],
    }
    # We return a Command object in the tool to update our state.
    return Command(update=state_update)

tools = [get_birthday, human_assistance]

llm_with_tools = llm.bind_tools(tools)

# chatbot node
def chatbot(state: State):
    message = llm_with_tools.invoke(state["messages"])
    # Because we will be interrupting during tool execution,
    # we disable parallel tool calling to avoid repeating any
    # tool invocations when we resume.
    assert len(message.tool_calls) <= 1
    return {"messages": [message]}

tool_node = ToolNode(tools=tools)

graph_builder = StateGraph(State)

# The first argument is the unique node name
# The second argument is the function or object that will be called whenever
# the node is used.
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_node("tools", tool_node)
# The `tools_condition` function returns "tools" if the chatbot asks to use a tool, and "END" if
# it is fine directly responding. This conditional routing defines the main agent loop.
graph_builder.add_conditional_edges(
    "chatbot",
    tools_condition,
)
# Any time a tool is called, we return to the chatbot to decide the next step
graph_builder.add_edge("tools", "chatbot")
graph_builder.add_edge(START, "chatbot")

graph = graph_builder.compile(checkpointer=memory)

my_config = {"configurable": {"thread_id": "1"}}

# human_response = (
#     "We, the experts are here to help! We'd recommend you check out LangGraph to build your agent."
#     " It's much more reliable and extensible than simple autonomous agents."
# )

human_command = Command(
    resume={
        "correct": "y",
        "name": "Hurry",
        "birthday": "May 05, 2028",
    },
)

# Can you tell me what is the birthday of Hurry? when you have the answer, use the human_assistance tool for review.

def stream_graph_updates(user_input: str):
    # print("+++++++++++ get state start +++++++++++")
    # snapshot = graph.get_state(my_config)
    # for k, v in snapshot.values.items():
    #     # if k in ("name", "birthday"):
    #     print(f"Key: {k}, Value: {v} \n")
    # print("+++++++++++ get state end +++++++++++")
    my_input = None
    if user_input.lower() == "human":
        my_input = human_command
    else:
        my_input = {"messages": [{"role": "user", "content": user_input}]}

    for event in graph.stream(input=my_input,
                              config=my_config,
                              stream_mode="values",
                              ):
        event["messages"][-1].pretty_print()
        # for value in event.values():
            # print(f"Assistant: {value[-1].content}")


while True:
    try:
        user_input = input("User: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break
    
        stream_graph_updates(user_input)
    except:
        # fallback if input() is not available
        # user_input = "What do you know about LangGraph?"
        # print("User: " + user_input)
        # stream_graph_updates(user_input)
        break

